from langchain_community.document_loaders.csv_loader import CSVLoader
import asyncio
import tempfile
from io import StringIO

file_path = r'D:\00_study\07_python\PythonProject\src' \
            r'\document_loaders\examples\mlb_teams_2012.csv'

async def load_csv():
    loader = CSVLoader(file_path=file_path)
    data = loader.load()
    for record in data[:2]:
        print(record)

async def load_customized_csv():
    loader = CSVLoader(file_path=file_path,
                       csv_args={
                           "delimiter": ",",
                           "quotechar": '"',
                           "fieldnames": ["MLB Team", "Payroll in millions", "Wins"],
                       },
                       )
    data = loader.load()
    for record in data[:2]:
        print(record)
async def load_specify_source():
    loader = CSVLoader(file_path=file_path,
                       source_column="Team"
                       )
    data = loader.load()
    for record in data[:2]:
        print(record)

async def load_with_tempfile():
    string_data = """
    "Team", "Payroll (millions)", "Wins"
    "Nationals",     81.34, 98
    "Reds",          82.20, 97
    "Yankees",      197.96, 95
    "Giants",       117.62, 94
    """.strip()

    with tempfile.NamedTemporaryFile(delete=False, mode="w+") as temp_file:
        temp_file.write(string_data)
        temp_file_path = temp_file.name
    loader = CSVLoader(file_path=temp_file_path,
                       )
    data = loader.load()
    for record in data[:2]:
        print(record)

asyncio.run(load_specify_source())
